import sys
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
                             QLabel, QPushButton, QGridLayout, QLineEdit)

class ArmController(QWidget):
    def __init__(self):
        super().__init__()
        self.controller = None
        self.jointAngles1 = [0.0] * 6
        self.jointAngles2 = [0.0] * 4
        self.speed = 50.0
        self.step = 5.0
        self.initUI()

    def initUI(self):
        self.setWindowTitle('Arm Controller')
        self.setGeometry(100, 100, 800, 500)

        layout = QVBoxLayout()

        # 连接按钮
        connectButton = QPushButton('Connect Controller', self)
        connectButton.clicked.connect(self.connectController)
        layout.addWidget(connectButton)

        # 参数设置
        param_layout = QHBoxLayout()
        # 速度设置
        speed_layout = QHBoxLayout()
        speedLabel = QLabel('Speed (mm/s):')
        self.speedLineEdit = QLineEdit(str(self.speed))
        self.speedLineEdit.setAlignment(Qt.AlignRight)
        self.speedLineEdit.textChanged.connect(self.updateSpeed)
        speed_layout.addWidget(speedLabel)
        speed_layout.addWidget(self.speedLineEdit)
        
        # 步长设置
        step_layout = QHBoxLayout()
        stepLabel = QLabel('Step Size:')
        self.stepLineEdit = QLineEdit(str(self.step))
        self.stepLineEdit.setAlignment(Qt.AlignRight)
        self.stepLineEdit.textChanged.connect(self.updateStep)
        step_layout.addWidget(stepLabel)
        step_layout.addWidget(self.stepLineEdit)
        
        param_layout.addLayout(speed_layout)
        param_layout.addLayout(step_layout)
        layout.addLayout(param_layout)

        # 机械臂布局
        armsLayout = QHBoxLayout()

        # 机械臂1（6个关节）
        arm1Layout = QVBoxLayout()
        arm1Layout.addWidget(QLabel('<b>Arm 1 (6 Joints)</b>', alignment=Qt.AlignCenter))
        
        grid1 = QGridLayout()
        self.valueLabels1 = []
        for i in range(6):
            # 关节标签
            joint_label = QLabel(f'Joint {i+1}:')
            # 当前值标签
            value_label = QLabel('0.0')
            value_label.setAlignment(Qt.AlignRight)
            value_label.setFixedWidth(60)
            self.valueLabels1.append(value_label)
            
            # 控制按钮
            minus_btn = QPushButton('-')
            minus_btn.setFixedWidth(40)
            minus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(1, idx, -1))
            
            plus_btn = QPushButton('+')
            plus_btn.setFixedWidth(40)
            plus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(1, idx, 1))
            
            grid1.addWidget(joint_label, i, 0)
            grid1.addWidget(value_label, i, 1)
            grid1.addWidget(minus_btn, i, 2)
            grid1.addWidget(plus_btn, i, 3)

        # 移动按钮
        moveButton1 = QPushButton('Move Arm 1', self)
        moveButton1.clicked.connect(lambda: self.moveArm(1))
        arm1Layout.addLayout(grid1)
        arm1Layout.addWidget(moveButton1)

        # 机械臂2（4个关节）
        arm2Layout = QVBoxLayout()
        arm2Layout.addWidget(QLabel('<b>Arm 2 (4 Joints)</b>', alignment=Qt.AlignCenter))
        
        grid2 = QGridLayout()
        self.valueLabels2 = []
        for i in range(4):
            joint_label = QLabel(f'Joint {i+7}:')
            value_label = QLabel('0.0')
            value_label.setAlignment(Qt.AlignRight)
            value_label.setFixedWidth(60)
            self.valueLabels2.append(value_label)
            
            minus_btn = QPushButton('-')
            minus_btn.setFixedWidth(40)
            minus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(2, idx, -1))
            
            plus_btn = QPushButton('+')
            plus_btn.setFixedWidth(40)
            plus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(2, idx, 1))
            
            grid2.addWidget(joint_label, i, 0)
            grid2.addWidget(value_label, i, 1)
            grid2.addWidget(minus_btn, i, 2)
            grid2.addWidget(plus_btn, i, 3)

        # 移动按钮
        moveButton2 = QPushButton('Move Arm 2', self)
        moveButton2.clicked.connect(lambda: self.moveArm(2))
        arm2Layout.addLayout(grid2)
        arm2Layout.addWidget(moveButton2)

        armsLayout.addLayout(arm1Layout)
        armsLayout.addLayout(arm2Layout)

        layout.addLayout(armsLayout)
        self.setLayout(layout)

    def connectController(self):
        if not self.controller:
            from driver_server import VSRController  # 假设VSRController可以从main.py导入
            try:
                self.controller = VSRController(simulate=True)
                self.controller.connect()
                print("Controller connected.")
            except Exception as e:
                print(f"Connection failed: {str(e)}")

    def adjust_joint_angle(self, arm_index, joint_index, direction):
        """调整关节角度值"""
        try:
            # 获取当前步长
            step = abs(float(self.stepLineEdit.text()))
        except ValueError:
            step = 5.0  # 默认步长
            
        delta = direction * step
        
        try:
            if arm_index == 1:
                current = self.jointAngles1[joint_index]
                new_value = max(-100.0, min(100.0, current + delta))
                self.jointAngles1[joint_index] = new_value
                self.valueLabels1[joint_index].setText(f"{new_value:.1f}")
            elif arm_index == 2:
                current = self.jointAngles2[joint_index]
                new_value = max(-100.0, min(100.0, current + delta))
                self.jointAngles2[joint_index] = new_value
                self.valueLabels2[joint_index].setText(f"{new_value:.1f}")
        except Exception as e:
            print(f"Error adjusting joint angle: {e}")

    def updateSpeed(self, text):
        """更新速度值"""
        try:
            self.speed = max(0.1, float(text))  # 最小速度0.1
        except ValueError:
            pass

    def updateStep(self, text):
        """更新步长值"""
        try:
            self.step = max(0.1, float(text))  # 最小步长0.1
        except ValueError:
            pass

    def moveArm(self, arm_index):
        """移动机械臂"""
        if not self.controller:
            print("Please connect the controller first.")
            return

        try:
            if arm_index == 1:
                target = self.jointAngles1
            else:
                target = self.jointAngles2
                
            print(f"Moving arm {arm_index} to: {target}")
            self.controller.set_joint_angle(target, self.speed)
        except Exception as e:
            print(f"Motion planning failed: {str(e)}")
            # 这里可以添加更详细的错误处理逻辑

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = ArmController()
    ex.show()
    sys.exit(app.exec_())

# import sys
# from PyQt5.QtCore import Qt
# from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
#                              QLabel, QPushButton, QGridLayout, QLineEdit)

# class ArmController(QWidget):
#     def __init__(self):
#         super().__init__()
#         self.controller = None
#         self.jointAngles1 = [0.0] * 6
#         self.jointAngles2 = [0.0] * 4
#         self.speed = 50.0
#         self.initUI()

#     def initUI(self):
#         self.setWindowTitle('Arm Controller')
#         self.setGeometry(100, 100, 800, 500)

#         layout = QVBoxLayout()

#         # 连接按钮
#         connectButton = QPushButton('Connect Controller', self)
#         connectButton.clicked.connect(self.connectController)
#         layout.addWidget(connectButton)

#         # 机械臂布局
#         armsLayout = QHBoxLayout()

#         # 机械臂1（6个关节）
#         arm1Layout = QVBoxLayout()
#         arm1Layout.addWidget(QLabel('<b>Arm 1 (6 Joints)</b>', alignment=Qt.AlignCenter))
        
#         grid1 = QGridLayout()
#         self.valueLabels1 = []
#         for i in range(6):
#             # 关节标签
#             joint_label = QLabel(f'Joint {i+1}:')
#             # 当前值标签
#             value_label = QLabel('0.0')
#             value_label.setAlignment(Qt.AlignRight)
#             value_label.setFixedWidth(60)
#             self.valueLabels1.append(value_label)
            
#             # 控制按钮
#             minus_btn = QPushButton('-')
#             minus_btn.setFixedWidth(40)
#             minus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(1, idx, -5))
            
#             plus_btn = QPushButton('+')
#             plus_btn.setFixedWidth(40)
#             plus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(1, idx, 5))
            
#             grid1.addWidget(joint_label, i, 0)
#             grid1.addWidget(value_label, i, 1)
#             grid1.addWidget(minus_btn, i, 2)
#             grid1.addWidget(plus_btn, i, 3)

#         # 移动按钮
#         moveButton1 = QPushButton('Move Arm 1', self)
#         moveButton1.clicked.connect(lambda: self.moveArm(1))
#         arm1Layout.addLayout(grid1)
#         arm1Layout.addWidget(moveButton1)

#         # 机械臂2（4个关节）
#         arm2Layout = QVBoxLayout()
#         arm2Layout.addWidget(QLabel('<b>Arm 2 (4 Joints)</b>', alignment=Qt.AlignCenter))
        
#         grid2 = QGridLayout()
#         self.valueLabels2 = []
#         for i in range(4):
#             joint_label = QLabel(f'Joint {i+7}:')
#             value_label = QLabel('0.0')
#             value_label.setAlignment(Qt.AlignRight)
#             value_label.setFixedWidth(60)
#             self.valueLabels2.append(value_label)
            
#             minus_btn = QPushButton('-')
#             minus_btn.setFixedWidth(40)
#             minus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(2, idx, -5))
            
#             plus_btn = QPushButton('+')
#             plus_btn.setFixedWidth(40)
#             plus_btn.clicked.connect(lambda _, idx=i: self.adjust_joint_angle(2, idx, 5))
            
#             grid2.addWidget(joint_label, i, 0)
#             grid2.addWidget(value_label, i, 1)
#             grid2.addWidget(minus_btn, i, 2)
#             grid2.addWidget(plus_btn, i, 3)

#         # 移动按钮
#         moveButton2 = QPushButton('Move Arm 2', self)
#         moveButton2.clicked.connect(lambda: self.moveArm(2))
#         arm2Layout.addLayout(grid2)
#         arm2Layout.addWidget(moveButton2)

#         armsLayout.addLayout(arm1Layout)
#         armsLayout.addLayout(arm2Layout)

#         # 速度设置
#         speedLayout = QHBoxLayout()
#         speedLabel = QLabel('Speed (mm/s):')
#         self.speedLineEdit = QLineEdit(str(self.speed))
#         self.speedLineEdit.setAlignment(Qt.AlignRight)
#         self.speedLineEdit.textChanged.connect(self.updateSpeed)
#         speedLayout.addWidget(speedLabel)
#         speedLayout.addWidget(self.speedLineEdit)

#         layout.addLayout(armsLayout)
#         layout.addLayout(speedLayout)
#         self.setLayout(layout)

#     def closeEvent(self, event):
#         """重写关闭事件处理"""
#         if self.controller:
#             self.controller.stop()
#             print("Controller disconnected.")
#         event.accept()


#     def connectController(self):
#         if not self.controller:
#             from main import VSRController  # 假设VSRController可以从main.py导入
#             self.controller = VSRController(simulate=True)
#             self.controller.connect()
#             print("Controller connected.")

#     def adjust_joint_angle(self, arm_index, joint_index, delta):
#         """调整关节角度值"""
#         try:
#             if arm_index == 1:
#                 current = self.jointAngles1[joint_index]
#                 new_value = max(-100.0, min(100.0, current + delta))
#                 self.jointAngles1[joint_index] = new_value
#                 self.valueLabels1[joint_index].setText(f"{new_value:.1f}")
#             elif arm_index == 2:
#                 current = self.jointAngles2[joint_index]
#                 new_value = max(-100.0, min(100.0, current + delta))
#                 self.jointAngles2[joint_index] = new_value
#                 self.valueLabels2[joint_index].setText(f"{new_value:.1f}")
#         except Exception as e:
#             print(f"Error adjusting joint angle: {e}")

#     def updateSpeed(self, text):
#         """更新速度值"""
#         try:
#             self.speed = float(text)
#         except ValueError:
#             pass

#     def moveArm(self, arm_index):
#         """移动机械臂"""
#         if not self.controller:
#             print("Please connect the controller first.")
#             return

#         if arm_index == 1:
#             target = self.jointAngles1 + self.jointAngles2
#         else:
#             target = self.jointAngles1 + self.jointAngles2
            
#         print(f"Moving arm {arm_index} to: {target}")
#         self.controller.set_joint_angle(target, self.speed)

# if __name__ == '__main__':
#     app = QApplication(sys.argv)
#     ex = ArmController()
#     ex.show()
#     sys.exit(app.exec_())